"""
Enhanced Subscription Example for vnpy_xt

This example demonstrates how to use the enhanced subscription features
that fully leverage xtdata.py capabilities:

1. Multi-period subscriptions (tick, 1m, 5m, 1d, etc.)
2. Dividend adjustment types (front/back/none)
3. L2 thousand-level market depth
4. Whole market quote subscription
"""

import sys
from pathlib import Path

# Add parent directories to path
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root / "vnpy"))
sys.path.insert(0, str(project_root / "vnpy_xt"))

from vnpy.event import EventEngine
from vnpy.trader.engine import MainEngine
from vnpy.trader.constant import Exchange

from vnpy_xt import (
    XtGateway,
    XtSubscribeRequest,
    XtL2SubscribeRequest,
    XtWholeQuoteRequest
)


def main():
    """示例主函数"""
    
    # 1. Initialize VNPy
    event_engine = EventEngine()
    main_engine = MainEngine(event_engine)
    main_engine.add_gateway(XtGateway)
    
    # 2. Connect to QMT
    xt_setting = {
        "token": "",
        "股票市场": "是",
        "期货市场": "否",
        "期权市场": "否",
        "仿真交易": "是",
        "账号类型": "股票",
        "QMT路径": "D:\\gjzqqmt_ceshi\\userdata_mini",
        "资金账号": "40743071"
    }
    
    main_engine.connect(xt_setting, "XT")
    
    # Get gateway
    gateway = main_engine.get_gateway("XT")
    
    if not gateway or not hasattr(gateway, 'md_api'):
        print("❌ Gateway not connected")
        return
    
    md_api = gateway.md_api
    
    print("\n" + "="*80)
    print("Enhanced Subscription Examples".center(80))
    print("="*80 + "\n")
    
    # ========================================================================
    # Example 1: Basic Tick Subscription (Compatible with old API)
    # ========================================================================
    print("📊 Example 1: Basic Tick Subscription")
    print("-" * 80)
    
    from vnpy.trader.object import SubscribeRequest
    
    tick_req = SubscribeRequest(
        symbol="600519",
        exchange=Exchange.SSE
    )
    md_api.subscribe(tick_req)
    print("✅ Subscribed to 600519.SSE tick data (compatible mode)\n")
    
    # ========================================================================
    # Example 2: 1-Minute Bar Subscription
    # ========================================================================
    print("📊 Example 2: 1-Minute Bar Subscription")
    print("-" * 80)
    
    min1_req = XtSubscribeRequest(
        symbol="600519",
        exchange=Exchange.SSE,
        period="1m"  # 1-minute bars
    )
    md_api.subscribe(min1_req)
    print("✅ Subscribed to 600519.SSE 1-minute bars\n")
    
    # ========================================================================
    # Example 3: Daily Bar with Forward Adjustment
    # ========================================================================
    print("📊 Example 3: Daily Bar with Forward Adjustment")
    print("-" * 80)
    
    daily_front_req = XtSubscribeRequest(
        symbol="600519",
        exchange=Exchange.SSE,
        period="1d",
        dividend_type="front"  # Forward adjustment (前复权)
    )
    md_api.subscribe(daily_front_req)
    print("✅ Subscribed to 600519.SSE daily bars with forward adjustment\n")
    
    # ========================================================================
    # Example 4: Historical Data with Count
    # ========================================================================
    print("📊 Example 4: Historical + Real-time Data")
    print("-" * 80)
    
    hist_req = XtSubscribeRequest(
        symbol="000001",
        exchange=Exchange.SZSE,
        period="5m",
        count=100,  # Get last 100 bars
        dividend_type="back"  # Backward adjustment (后复权)
    )
    md_api.subscribe(hist_req)
    print("✅ Subscribed to 000001.SZSE 5-minute bars (last 100) with backward adjustment\n")
    
    # ========================================================================
    # Example 5: Time Range Subscription
    # ========================================================================
    print("📊 Example 5: Time Range Subscription")
    print("-" * 80)
    
    range_req = XtSubscribeRequest(
        symbol="600000",
        exchange=Exchange.SSE,
        period="1d",
        start_time="20240101",
        end_time="20241231",
        dividend_type="none"  # No adjustment (不复权)
    )
    md_api.subscribe(range_req)
    print("✅ Subscribed to 600000.SSE daily bars from 2024-01-01 to 2024-12-31\n")
    
    # ========================================================================
    # Example 6: L2 Thousand-Level Market Depth
    # ========================================================================
    print("📊 Example 6: L2 Thousand-Level Market Depth")
    print("-" * 80)
    
    def on_l2_data(data):
        """L2 data callback"""
        for symbol, quotes in data.items():
            if quotes:
                quote = quotes[0]
                print(f"L2更新 {symbol}: "
                      f"买{len(quote.get('bidPrice', []))}档, "
                      f"卖{len(quote.get('askPrice', []))}档")
    
    l2_req = XtL2SubscribeRequest(
        symbol="600519",
        exchange=Exchange.SSE,
        gear_num=10,  # Subscribe to 10 gears
        callback=on_l2_data
    )
    seq = md_api.subscribe_l2_thousand(l2_req)
    
    if seq > 0:
        print(f"✅ Subscribed to 600519.SSE L2 thousand-level depth (10 gears), seq={seq}\n")
    else:
        print(f"⚠️  L2 subscription failed (seq={seq}). This is normal if you don't have L2 permission.\n")
        print("   To get L2 data, you need to:")
        print("   1. Apply for Level-2 market data permission from your broker")
        print("   2. Ensure your QMT account has L2 subscription activated\n")
    
    # ========================================================================
    # Example 7: Whole Market Quote (Entire Exchange)
    # ========================================================================
    print("📊 Example 7: Whole Market Quote Subscription")
    print("-" * 80)
    
    def on_whole_quote(data):
        """Whole quote callback"""
        from datetime import datetime
        
        total_count = len(data)
        current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        print(f"\n📈 全推行情更新 [{current_time}]: 收到 {total_count} 个合约数据")
        
        if total_count == 0:
            return
        
        # Display data in DataFrame-style format (first 2 + last 2 rows)
        print("=" * 120)
        
        # Header
        header = f"{'股票代码':<12} | {'最新价':>10} | {'开盘':>10} | {'最高':>10} | {'最低':>10} | {'成交量':>12} | {'成交额':>14}"
        print(header)
        print("-" * 120)
        
        # Get items as list
        items = list(data.items())
        
        # Determine which rows to display
        if total_count <= 4:
            display_items = items
        else:
            # Show first 2 and last 2 rows
            display_items = items[:2] + [('...', None)] + items[-2:]
        
        # Display rows
        for code, quote_data in display_items:
            if code == '...':
                print(f"{'...':<12} | {'...':>10} | {'...':>10} | {'...':>10} | {'...':>10} | {'...':>12} | {'...':>14}")
                continue
            
            if not quote_data:
                continue
            
            # Handle different data structures
            if isinstance(quote_data, list) and quote_data:
                quote = quote_data[0]  # Take first item if list
            else:
                quote = quote_data
            
            # Extract values with fallbacks
            if isinstance(quote, dict):
                last_price = quote.get('lastPrice', quote.get('last_price', 0))
                open_price = quote.get('open', 0)
                high_price = quote.get('high', 0)
                low_price = quote.get('low', 0)
                volume = quote.get('volume', 0)
                amount = quote.get('amount', 0)
            else:
                # Handle object attributes
                last_price = getattr(quote, 'lastPrice', getattr(quote, 'last_price', 0))
                open_price = getattr(quote, 'open', 0)
                high_price = getattr(quote, 'high', 0)
                low_price = getattr(quote, 'low', 0)
                volume = getattr(quote, 'volume', 0)
                amount = getattr(quote, 'amount', 0)
            
            # Format row
            row = (
                f"{code:<12} | "
                f"{last_price:>10.2f} | "
                f"{open_price:>10.2f} | "
                f"{high_price:>10.2f} | "
                f"{low_price:>10.2f} | "
                f"{volume:>12,} | "
                f"{amount:>14,.0f}"
            )
            print(row)
        
        # Footer
        print("=" * 120)
        print(f"[{total_count} rows]\n")
    
    whole_req = XtWholeQuoteRequest(
        code_list=["SH", "SZ"],  # Subscribe to Shanghai and Shenzhen markets
        callback=on_whole_quote
    )
    seq = md_api.subscribe_whole_quote(whole_req)
    print(f"✅ Subscribed to whole market quote (SH + SZ), seq={seq}\n")
    
    # ========================================================================
    # Example 8: Multiple Stocks Whole Quote
    # ========================================================================
    print("📊 Example 8: Specific Stocks Whole Quote")
    print("-" * 80)
    
    stocks_req = XtWholeQuoteRequest(
        code_list=["600519.SH", "000001.SZ", "600000.SH"]
    )
    seq = md_api.subscribe_whole_quote(stocks_req)
    print(f"✅ Subscribed to specific stocks whole quote, seq={seq}\n")
    
    print("="*80)
    print("All subscriptions completed!".center(80))
    print("Press Ctrl+C to exit".center(80))
    print("="*80 + "\n")
    
    # Keep running
    try:
        import time
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n👋 Shutting down...")
        main_engine.close()


if __name__ == "__main__":
    main()
